#include "myHiGMM.h"
#include "parameter.h"

using namespace cv;

MyHiGMM::MyHiGMM()
{
  memset(&s_stGmmInfo,0,sizeof(s_stGmmInfo));

  //s_stGmmSwitch.bVenc = HI_TRUE;
  //s_stGmmSwitch.bVo   = HI_TRUE;

  //s_stViConfig = {0};
}

void MyHiGMM::Set_Learning_Rate(int learn_rate)
{
  this->s_stGmmInfo.stGmm.stGmmCtrl.u0q16LearnRate = learn_rate;
}

void MyHiGMM::Set_Frame_Width_Height(int width, int height)
{
  this->frame_width = width;
  this->frame_height = height;
}

HI_VOID MyHiGMM::SAMPLE_IVE_Gmm_Uninit(SAMPLE_IVE_GMM_S* pstGmm)
{
  IVE_MMZ_FREE(pstGmm->stSrc.au64PhyAddr[0], pstGmm->stSrc.au64VirAddr[0]);
  IVE_MMZ_FREE(pstGmm->stFg.au64PhyAddr[0], pstGmm->stFg.au64VirAddr[0]);
  IVE_MMZ_FREE(pstGmm->stBg.au64PhyAddr[0], pstGmm->stBg.au64VirAddr[0]);
  IVE_MMZ_FREE(pstGmm->stModel.u64PhyAddr, pstGmm->stModel.u64VirAddr);
  IVE_MMZ_FREE(pstGmm->stImg1.au64PhyAddr[0], pstGmm->stImg1.au64VirAddr[0]);
  IVE_MMZ_FREE(pstGmm->stImg2.au64PhyAddr[0], pstGmm->stImg2.au64VirAddr[0]);
  IVE_MMZ_FREE(pstGmm->stBlob.u64PhyAddr, pstGmm->stBlob.u64VirAddr);
}

HI_S32 MyHiGMM::SAMPLE_IVE_Gmm_Init(SAMPLE_IVE_GMM_S* pstGmm, HI_U32 u32Width, HI_U32 u32Height)
{
    //this->Set_Frame_Width_Height((int)u32Width, (int)u32Height);

    HI_S32 s32Ret = HI_SUCCESS;
    HI_U32 u32Size = 0;
    HI_S8 as8Mask[25] = {1, 2, 3, 2, 1,
                         2, 5, 6, 5, 2,
                         3, 6, 8, 6, 3,
                         2, 5, 6, 5, 2,
                         1, 2, 3, 2, 1
                        };

    memset(pstGmm, 0, sizeof(SAMPLE_IVE_GMM_S));

    s32Ret = SAMPLE_COMM_IVE_CreateImage(&yuv420Img, IVE_IMAGE_TYPE_YUV420SP, u32Width, u32Height); // for U8C3
        SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret, s32Ret, "Error(%#x),Create yuv420Img image failed!\n", s32Ret);

    //s32Ret = SAMPLE_COMM_IVE_CreateImage(&(pstGmm->stSrc), IVE_IMAGE_TYPE_U8C1, u32Width, u32Height); // for U8C1
    s32Ret = SAMPLE_COMM_IVE_CreateImage(&(pstGmm->stSrc), IVE_IMAGE_TYPE_U8C3_PACKAGE, u32Width, u32Height); // for U8C3
    SAMPLE_CHECK_EXPR_GOTO(HI_SUCCESS != s32Ret, GMM_INIT_FAIL,
        "Error(%#x),Create src image failed!\n", s32Ret);

    s32Ret = SAMPLE_COMM_IVE_CreateImage(&(pstGmm->stFg), IVE_IMAGE_TYPE_U8C1, u32Width, u32Height);
    SAMPLE_CHECK_EXPR_GOTO(HI_SUCCESS != s32Ret, GMM_INIT_FAIL,
        "Error(%#x),Create fg image failed!\n", s32Ret);

    //s32Ret = SAMPLE_COMM_IVE_CreateImage(&(pstGmm->stBg), IVE_IMAGE_TYPE_U8C1, u32Width, u32Height); // for U8C1
    s32Ret = SAMPLE_COMM_IVE_CreateImage(&(pstGmm->stBg), IVE_IMAGE_TYPE_U8C3_PACKAGE, u32Width, u32Height); // for U8C3
    SAMPLE_CHECK_EXPR_GOTO(HI_SUCCESS != s32Ret, GMM_INIT_FAIL,
        "Error(%#x),Create bg image failed!\n", s32Ret);

    pstGmm->stGmmCtrl.u0q16InitWeight = 3276; //0.05
    pstGmm->stGmmCtrl.u0q16BgRatio = 45875;   //0.8
    pstGmm->stGmmCtrl.u22q10MaxVar = 3 * 2000 * 1024;//(2000 << 10);
    pstGmm->stGmmCtrl.u22q10MinVar = 3 * 200 * 1024;//(200 << 10);
    pstGmm->stGmmCtrl.u22q10NoiseVar = 3 * 225 * 1024;//(225 << 10);
    pstGmm->stGmmCtrl.u8q8VarThr = (HI_U16)(256 * 6.25);//1600;
    pstGmm->stGmmCtrl.u8ModelNum = 3;
    //pstGmm->stGmmCtrl.u0q16LearnRate = 327;

    //u32Size = pstGmm->stSrc.u32Width * pstGmm->stSrc.u32Height * pstGmm->stGmmCtrl.u8ModelNum * 7; //for U8C1
    u32Size = pstGmm->stSrc.u32Width * pstGmm->stSrc.u32Height * pstGmm->stGmmCtrl.u8ModelNum * 11; //for U8C3
    s32Ret = SAMPLE_COMM_IVE_CreateMemInfo(&pstGmm->stModel, u32Size);
    SAMPLE_CHECK_EXPR_GOTO(HI_SUCCESS != s32Ret, GMM_INIT_FAIL,
        "Error(%#x),Create model mem info failed!\n", s32Ret);
    memset(SAMPLE_COMM_IVE_CONVERT_64BIT_ADDR(HI_VOID,pstGmm->stModel.u64VirAddr), 0, pstGmm->stModel.u32Size);

    s32Ret = SAMPLE_COMM_IVE_CreateImage(&(pstGmm->stImg1), IVE_IMAGE_TYPE_U8C1, u32Width, u32Height);
      SAMPLE_CHECK_EXPR_GOTO(HI_SUCCESS != s32Ret, GMM_INIT_FAIL,
        "Error(%#x),Create img1 image failed!\n", s32Ret);

    s32Ret = SAMPLE_COMM_IVE_CreateImage(&(pstGmm->stImg2), IVE_IMAGE_TYPE_U8C1, u32Width, u32Height);
      SAMPLE_CHECK_EXPR_GOTO(HI_SUCCESS != s32Ret, GMM_INIT_FAIL,
        "Error(%#x),Create img2 image failed!\n", s32Ret);

    u32Size = sizeof(IVE_CCBLOB_S);
    s32Ret = SAMPLE_COMM_IVE_CreateMemInfo(&pstGmm->stBlob, u32Size);
      SAMPLE_CHECK_EXPR_GOTO(HI_SUCCESS != s32Ret, GMM_INIT_FAIL,
        "Error(%#x),Create blob mem info failed!\n", s32Ret);

    memcpy(pstGmm->stFltCtrl.as8Mask, as8Mask, 25);
    pstGmm->stFltCtrl.u8Norm = 7;
    memset(pstGmm->stDilateCtrl.au8Mask, 255, 25);
    memset(pstGmm->stErodeCtrl.au8Mask, 255, 25);
    pstGmm->stCclCtrl.u16InitAreaThr     = 16;
    pstGmm->stCclCtrl.u16Step             = 4;
    pstGmm->stCclCtrl.enMode             = IVE_CCL_MODE_8C;

GMM_INIT_FAIL:

    if (HI_SUCCESS != s32Ret)
    {
        SAMPLE_IVE_Gmm_Uninit(pstGmm);
    }
    return s32Ret;
}

HI_S32 MyHiGMM::SAMPLE_IVE_GmmProc(SAMPLE_IVE_GMM_INFO_S* pstGmmInfo, VIDEO_FRAME_INFO_S* pstExtFrmInfo,HI_U32 u32BaseWidth,HI_U32 u32BaseHeight)
{
  HI_S32 s32Ret;
    SAMPLE_IVE_GMM_S* pstGmm;
    IVE_HANDLE IveHandle;
    HI_BOOL bFinish = HI_FALSE;
    HI_BOOL bBlock = HI_TRUE;
    HI_BOOL bInstant = HI_TRUE;
    IVE_CCBLOB_S* pstBlob;
    IVE_CSC_CTRL_S pstCscCtrl;

    pstGmm = &pstGmmInfo->stGmm;
    pstBlob = SAMPLE_COMM_IVE_CONVERT_64BIT_ADDR(IVE_CCBLOB_S,pstGmm->stBlob.u64VirAddr);

    //get YUV2RGB ive img
    bInstant = HI_FALSE;
    s32Ret = SAMPLE_COMM_IVE_DmaImage(pstExtFrmInfo,&yuv420Img,bInstant);
    SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),SAMPLE_COMM_IVE_DmaImage failed!\n",s32Ret);
    pstCscCtrl.enMode = IVE_CSC_MODE_PIC_BT601_YUV2RGB;
    s32Ret = HI_MPI_IVE_CSC(&IveHandle, &yuv420Img, &pstGmm->stSrc, &pstCscCtrl, bInstant);
    SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_CSC failed!\n",s32Ret);

    //1.Get Y
    //bInstant = HI_FALSE;
    //s32Ret = SAMPLE_COMM_IVE_DmaImage(pstExtFrmInfo,&pstGmm->stSrc,bInstant);
    //SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),SAMPLE_COMM_IVE_DmaImage failed!\n",s32Ret);

    //2.Filter
    //s32Ret = HI_MPI_IVE_Filter(&IveHandle, &pstGmm->stSrc, &pstGmm->stImg1, &pstGmm->stFltCtrl, bInstant);
    //SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_Filter failed!\n",s32Ret);

    //3.Gmm
    s32Ret = HI_MPI_IVE_GMM(&IveHandle, &pstGmm->stSrc, &pstGmm->stFg, &pstGmm->stBg, &pstGmm->stModel, &pstGmm->stGmmCtrl, bInstant);
     SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_GMM failed!\n",s32Ret);

    //4.Dilate
    s32Ret = HI_MPI_IVE_Dilate(&IveHandle, &pstGmm->stFg, &pstGmm->stImg1, &pstGmm->stDilateCtrl, bInstant);
     SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_Dilate failed!\n",s32Ret);

    //5.Erode
    s32Ret = HI_MPI_IVE_Erode(&IveHandle, &pstGmm->stImg1, &pstGmm->stImg2, &pstGmm->stErodeCtrl, bInstant);
    SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_Erode failed!\n",s32Ret);

    //6.CCL
    bInstant = HI_TRUE;
    s32Ret = HI_MPI_IVE_CCL(&IveHandle, &pstGmm->stImg2, &pstGmm->stBlob, &pstGmm->stCclCtrl, bInstant);
    SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_CCL failed!\n",s32Ret);

    //Wait task finish
    s32Ret = HI_MPI_IVE_Query(IveHandle, &bFinish, bBlock);
    while (HI_ERR_IVE_QUERY_TIMEOUT == s32Ret)
    {
        usleep(100);
        s32Ret = HI_MPI_IVE_Query(IveHandle, &bFinish, bBlock);
    }
    SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_Query failed!\n",s32Ret);

    SAMPLE_COMM_IVE_BlobToRect(pstBlob, &(pstGmm->stRegion), IVE_RECT_NUM, 8,
        pstExtFrmInfo->stVFrame.u32Width, pstExtFrmInfo->stVFrame.u32Height,
        u32BaseWidth,u32BaseHeight);

    return HI_SUCCESS;
}

// only support CVBGR2IVEU8C3PACKAGE and CVGRAY2IVEU8C1
void MyHiGMM::mat2IveImg(Mat *src, IVE_IMAGE_S *pstDst, int type)
{
	if(src->cols != pstDst->u32Width || src->rows != pstDst->u32Height)
	{
		printf("width or height didn't equal, line:%d\n",__LINE__);
	}

	switch(type)
	{
	case CVGRAY_IVEU8C1:
		{
			for(int r=0; r<src->rows; r++)
			{
				HI_U8 *ptm = src->ptr(r);
				HI_U8 *pti = (HI_U8*)(pstDst->au64PhyAddr[0] + r * pstDst->au32Stride[0]);

				memcpy(pti, ptm, pstDst->u32Width);
			}
			break;
		}
	case CVBGR_IVEU8C3PACKAGE:
		{
			for(int r=0; r<src->rows; r++)
			{
				HI_U8 *ptm = src->ptr(r);
				HI_U8 *pti = (HI_U8*)(pstDst->au64PhyAddr[0] + 3 * r * pstDst->au32Stride[0]);

				memcpy(pti, ptm, pstDst->u32Width * 3);
			}
			break;
		}
	default:
		printf("not support type for mat2IveImg!\n");
	}
}

// only support IVEU8C1_CVGRAY and IVEU8C3PACKAGE_CVBGR
void MyHiGMM::iveImage2Mat(IVE_IMAGE_S *pstSrc, Mat *dst, int type)
{
	if(dst->cols != pstSrc->u32Width || dst->rows != pstSrc->u32Height)
	{
		printf("width or height didn't equal,line:%d\n",__LINE__);
		return;
	}

	switch(type)
	{
	case IVEU8C1_CVGRAY:
		{
			for(int r=0; r<dst->rows; r++)
			{
				HI_U8 *ptm = dst->ptr(r);
				HI_U8 *pti = (HI_U8*)(pstSrc->au64VirAddr[0] + r * pstSrc->au32Stride[0]);

				memcpy(ptm, pti, dst->cols);
			}
			break;
		}
	case IVEU8C3PACKAGE_CVBGR:
		{
			for(int r=0; r<dst->rows; r++)
			{
				HI_U8 *ptm = dst->ptr(r);
				HI_U8 *pti = (HI_U8*)(pstSrc->au64PhyAddr[0] + 3 * r * pstSrc->au32Stride[0]);
				//printf("width: %d,  stride: %d,  cols: %d,  height: %d,  rows: %d\n", pstSrc->u32Width, pstSrc->au32Stride[0], dst->cols, pstSrc->u32Height, dst->rows);
				printf("before memcpy\n");
				memcpy(ptm, pti, dst->cols * 3);
				/*HI_S32 s32Ret;
        IVE_HANDLE hIveHandle;
        IVE_SRC_DATA_S stSrcData;
        IVE_DST_DATA_S stDstData;
        IVE_DMA_CTRL_S stCtrl = {IVE_DMA_MODE_DIRECT_COPY,0};
        HI_BOOL bFinish = HI_FALSE;
        HI_BOOL bBlock = HI_TRUE;
        HI_BOOL bInstant = HI_FALSE;

        //fill src
        stSrcData.u64VirAddr = pstSrc->au64VirAddr[0];
        stSrcData.u64PhyAddr = pstSrc->au64PhyAddr[0];
        stSrcData.u32Width   = pstSrc->u32Width;
        stSrcData.u32Height  = pstSrc->u32Height;
        stSrcData.u32Stride  = pstSrc->au32Stride[0];

        //fill dst
        stDstData.u64VirAddr = dst->ptr(0);
        stDstData.u64PhyAddr = dst->ptr(0);
        stDstData.u32Width   = dst->cols;
        stDstData.u32Height  = dst->rows;
        stDstData.u32Stride  = dst->cols;

        s32Ret = HI_MPI_IVE_DMA(&hIveHandle,&stSrcData,&stDstData,&stCtrl,bInstant);
        SAMPLE_CHECK_EXPR_RET(HI_SUCCESS != s32Ret,s32Ret,"Error(%#x),HI_MPI_IVE_DMA failed!\n",s32Ret);
        */
        printf("after memcpy\n");
			}
			break;
		}
	default:
		printf("not support type for iveImage2Mat!\n");
	}
}
